Repository navigation
Add Databricks agent invocation operator and hook - #74386
Open
AlejandroMorgante wants to merge 10 commits into
Open
AlejandroMorgante wants to merge 10 commits into
AlejandroMorgante wants to merge 10 commits into
Conversation
Contributor
Author
|
@SameerMesiah97 could you please help me with a review when you have some time? Thanks! 🙏 |
AlejandroMorgante
marked this pull request as draft
October 9, 2026 19:31
AlejandroMorgante
marked this pull request as ready for review
October 9, 2026 19:31
kaxil
reviewed
Oct 9, 2026
kaxil
reviewed
Oct 9, 2026
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 9, 2026 23:25
152283d to
e13d7dd
Compare
kaxil
reviewed
Oct 10, 2026
AlejandroMorgante
marked this pull request as draft
October 10, 2026 15:25
Allow Airflow tasks to invoke DurableAgentServer agents deployed on Databricks Apps and wait without occupying a worker. Generated-by: Codex (GPT-6)
Keep the new operator documentation valid for Sphinx and recognize the official MLflow spelling. Generated-by: Codex (GPT-6)
Keep the operator guide focused on user configuration and invocation, following the existing Databricks documentation conventions. Generated-by: Codex (GPT-6)
A completed response must not turn an expired invocation wait into success. Recoverable async disconnects should not fail a task immediately, and polling failures need safe context for diagnosis. Generated-by: Codex (GPT-6)
Common AI consumers need to consult Databricks agents through the same contract as Vertex AI and Bedrock. Existing background and deferrable workflows must remain usable without installing the optional Common AI provider. Generated-by: Codex (GPT-6)
Replacement mocks should follow Airflow's testing conventions so they remain tied to the API they represent. Generated-by: Codex (GPT-6)
Mocked responses must reflect the Databricks runtime contract so session, interruption, failure and idempotency defects cannot hide behind passing tests. Non-default settings and unconditional assertions make routing and polling regressions observable. This draft PR step intentionally exposes failures pending implementation fixes in the following commits. Generated-by: Codex (GPT-6)
AlejandroMorgante
force-pushed
the
add-databricks-agent-invoke
branch
from
October 10, 2026 15:28
7266b14 to
7f3b704
Compare
Stored failures and template interruptions must not be treated as transient errors or completed answers. Task retries need to respect the server's idempotency rules, and toolset calls need a session. Generated-by: Codex (GPT-6)
Conditional expectations made individual regression cases harder to assess during review. Generated-by: Codex (GPT-6)
Make agent task progress traceable from submission through completion without exposing request or response contents. Generated-by: Codex (GPT-6)
Contributor
Author
End-to-end validationValidated commit Both synchronous and deferrable invocations completed successfully. Verification tasks confirmed the response content, original input, and session ID. The deferrable task resumed successfully after the trigger fired. Environment: Airflow 3.3.0, Python 3.12, PostgreSQL, and LocalExecutor. The test loaded the provider source from the PR commit. No model endpoint was used. Successful Airflow run
Synchronous invocation
Deferrable invocation
Agent input and output
Databricks App
Databricks invocation logs
The temporary OAuth secret was revoked and the Databricks App stopped after validation. |
AlejandroMorgante
marked this pull request as ready for review
October 11, 2026 01:05
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add
DatabricksAgentInvokeOperatorandDatabricksAgentHookto invoke agents hosted withDurableAgentServeron Databricks Apps. Supports service principal OAuth, background submission, synchronous polling and deferrable waiting, preserving the full invocation response in XCom.The hook implements
BaseManagedAgentHookfor Common AI and toolset use, with a default session per invocation. Request, state and error handling were checked against the runtime and templates in Databricks'databricks-ai-bridgerepository, tagdatabricks-agentbricks-v0.4.0. Stored failures are distinguished from transient HTTP errors, and interruptions nested in the output are returned by the operator and rejected by Common AI.Invocation IDs remain stable across retries and clears with unchanged input and session. Stored failures stop task retries and require a new ID to run again. Polling requests and retries respect the remaining timeout budget.
Includes documentation, focused unit tests and a deterministic system-test fixture requiring no model endpoint.
Validation of the latest changes:
tenacity==8.3.0.tomllib.Earlier validation: the full Databricks suite passed (1087 passed, 12 skipped), and the system test passed against a real workspace with OAuth, including deferral and resumption. Those broader runs have not been repeated after the latest changes.
See the fixture README for deployment, configuration and cleanup. The fixture uses an in-memory store and does not test recovery across app restarts. With the app and configuration prepared:
SYSTEM_TESTS_ENV_ID=your-unique-id \ BREEZE_INIT_COMMAND='set -a; . /files/databricks-agent-system-test.env; set +a' \ breeze testing system-tests \ --backend sqlite \ --forward-credentials \ --test-timeout 2400 \ providers/databricks/tests/system/databricks/example_databricks_agent.py \ -qWas generative AI tooling used to co-author this PR?
Generated-by: Codex (GPT-6) following the guidelines